# -*- coding: utf-8 -*-
"""
2.4 索引（Indexes）与检索（Retrieval） - Python 代码示例

本示例将演示构建一个简单 RAG (Retrieval-Augmented Generation) 流程的核心步骤：
1.  **准备数据:** 创建一个临时文本文件。
2.  **加载 (Load):** 使用 TextLoader 加载文本文件。
3.  **分割 (Split):** 使用 RecursiveCharacterTextSplitter 将文档分割成块。
4.  **嵌入 (Embed):** 使用 OpenAIEmbeddings 为文本块创建向量嵌入。
5.  **存储 (Store):** 使用 Chroma (内存向量数据库) 存储文本块和嵌入。
6.  **检索 (Retrieve):** 使用 VectorStore 创建 Retriever 并执行相似度搜索。

运行前准备：
1.  确保已安装 langchain, langchain-openai, python-dotenv, chromadb, tiktoken:
    `pip install langchain langchain-openai python-dotenv chromadb tiktoken`
    (tiktoken 通常是 openai 库的依赖，但显式安装更保险)
2.  确保项目根目录下或当前目录有 `.env` 文件，并包含 `OPENAI_API_KEY`。
3.  激活 Python 虚拟环境。
4.  运行: `python 2.4_Indexes_Retrieval.py`
"""

import os
import shutil
from dotenv import load_dotenv

# 1. 加载器
from langchain_community.document_loaders import TextLoader

# 2. 分割器
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 3. 嵌入模型
from langchain_huggingface import HuggingFaceEmbeddings


# 4. 向量存储
from langchain_community.vectorstores import Chroma

# 加载环境变量
load_dotenv()
print("环境变量已加载。")

# --- 1. 准备示例数据 --- 


# --- 2. 加载文档 --- 
print("\n--- 2. 加载文档 --- ")
DEMO_DIR = "./temp_data_for_indexing"
DEMO_FILE = os.path.join(DEMO_DIR, "sample_document.txt")
loader = TextLoader(DEMO_FILE, encoding="utf-8")
docs = loader.load()
print(f"文档加载完成。共加载 {len(docs)} 个文档。")
print(f"第一个文档内容预览 (前 50 字符): {docs[0].page_content[:50]}...")
print(f"第一个文档元数据: {docs[0].metadata}")

# --- 3. 分割文档 --- 
print("\n--- 3. 分割文档 --- ")
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=100,  # 每个块的目标大小 (字符数)
    chunk_overlap=20,  # 相邻块的重叠大小
    length_function=len,
    is_separator_regex=False,
)
splitted_docs = text_splitter.split_documents(docs)
print(f"文档分割完成。共分割成 {len(splitted_docs)} 个文本块。")
print("前三个文本块内容预览:")
for i, doc in enumerate(splitted_docs[:3]):
    print(f"  块 {i+1}: {doc.page_content}")

# --- 4. 初始化嵌入模型 --- 
print("\n--- 4. 初始化嵌入模型 --- ")
# 使用 OpenAI 的嵌入模型
# 请确保 OPENAI_API_KEY 已在 .env 或环境变量中设置
# 也可以选择本地模型，如 HuggingFaceEmbeddings
# embeddings = OpenAIEmbeddings(model="text-embedding-3-small") # 注释掉 OpenAI 初始化

# 使用从镜像站下载到本地的 HuggingFace 模型 BAAI/bge-large-zh-v1.5
# 请确保下面的路径指向你实际存放模型文件的文件夹
local_model_path = "./models/bge-large-zh-v1.5" # <--- 修改这里！指向包含模型文件的本地路径
model_kwargs = {'device': 'cpu'} # 如果有 CUDA 环境且安装了 torch GPU 版本, 可以改为 'cuda'
encode_kwargs = {'normalize_embeddings': True} # BGE 模型通常推荐归一化
embeddings = HuggingFaceEmbeddings(
    model_name=local_model_path, # <--- 使用本地路径加载模型
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs,
    # query_instruction="为这个句子生成表示以用于检索相关文章：" # BGE v1.5 版本对于 query 默认有指令，可以不用显式指定
)
print(f"嵌入模型 ({local_model_path}) 初始化成功。") # <--- 更新打印信息

# --- 5. 创建并填充向量存储 --- 
print("\n--- 5. 创建并填充向量存储 (Chroma) --- ")
# 使用 Chroma 作为内存向量数据库
# 它会自动处理文本块的嵌入过程
# persist_directory 参数可以指定持久化路径，如果省略则只在内存中
vectorstore = Chroma.from_documents(
    documents=splitted_docs,      # 传入分割后的文档块
    embedding=embeddings,       # 传入 HuggingFaceEmbeddings 实例
    # persist_directory="./chroma_db" # 可选：指定持久化目录
)
print("向量存储创建并填充完成 (使用 Chroma)。")
print(f"向量库中的文档数量: {vectorstore._collection.count()}") # 访问底层 collection 确认

# --- 6. 从向量存储创建检索器并检索 --- 
print("\n--- 6. 创建检索器并进行检索 --- ")
# 从 VectorStore 创建 Retriever
# k=2 表示返回最相似的 2 个结果
retriever = vectorstore.as_retriever(search_kwargs={"k": 2})
print("检索器创建成功。")

query = "LangChain 有哪些核心组件？"
print(f"\n使用查询进行检索: '{query}'")

# 使用 Retriever 进行检索
retrieved_docs = retriever.invoke(query)

print(f"检索到 {len(retrieved_docs)} 个相关文档块:")
for i, doc in enumerate(retrieved_docs):
    print(f"\n--- 相关文档 {i+1} ---")
    print(f"内容: {doc.page_content}")
    print(f"来源: {doc.metadata.get('source', '未知')}")

# --- 清理临时文件 --- 
# print("\n--- 清理临时文件 ---")
# try:
#     shutil.rmtree(DEMO_DIR)
#     # 如果指定了 persist_directory="./chroma_db"，也可以在这里清理
#     # if os.path.exists("./chroma_db"): shutil.rmtree("./chroma_db") 
#     print("临时文件和目录已清理。")
# except Exception as e:
#     print(f"清理临时文件出错: {e}")

print("\n索引与检索流程演示完毕。")
